#include <mm/pgtable.h>
#include <mm/mm.h>

// Provided by linker (see linker.ld)
extern uptr __kernel_start;
extern uptr __kernel_end;
extern uptr __init_hhk_srart;
extern uptr __init_hhk_end;
extern uptr __kernel_size_from_start;

#define sym_val(sym) (uptr)(&sym)

#define KERNEL_PAGE_COUNT                                                      \
    ((sym_val(__kernel_end) - sym_val(__kernel_start) + 0x1000 - 1) >> 12);
#define HHK_PAGE_COUNT ((sym_val(__init_hhk_end) - 0x100000 + 0x1000 - 1) >> 12)

void ptd_init(u32 ptd) {
    pg_table_t* dir = (pg_table_t*)ptd;

    dir->entry[0].data = ((uptr)dir + 0x1000) & DIR_BASE_VADDR;
    dir->entry[0].present = 1;
    dir->entry[0].rw = 1;
    dir->entry[0].user = 0;

    // 对低 1MiB 空间进行对等映射也包括了我们的VGA，方便内核操作。
    pg_table_t* tbl = (pg_table_t*)(dir->entry[0].data & DIR_BASE_VADDR);
    for (u32 i = 0; i < 256; i++) {
        tbl->entry[i].frame = i;
        tbl->entry[i].present = 1;
        tbl->entry[i].rw = 1;
        tbl->entry[i].user = 0;
    }

    // 对等映射我们的hhk_init，这样一来，当分页与地址转换开启后，我们依然能够照常执行最终的
    // jmp 指令来跳转至内核的入口点
    for (u32 i = 0; i < HHK_PAGE_COUNT; i++) {
        u32 index = i + 256; // 跳过前256个页(1MB)的对等映射
        tbl->entry[index].frame = index;
        tbl->entry[index].present = 1;
        tbl->entry[index].rw = 1;
        tbl->entry[index].user = 0;
    }

    // 将内核重映射至高半区
    u32 kernel_dir_index = DIRECTORY_INDEX(sym_val(__kernel_start));
    u32 kernel_tbl_index = TABLE_INDEX(sym_val(__kernel_start));
    u32 kernel_page_count = KERNEL_PAGE_COUNT;

    // 计算内核.text段的物理地址
    uptr kernel_pm = V2P(sym_val(__kernel_start));

    // 将内核所需要的页表注册进页目录
    for (u32 i = 0; i < KERNEL_MAX_TABLE; i++) {
        u32 index = i + kernel_dir_index;
        dir->entry[index].data = (ptd + 0x2000 + i * 0x1000) & DIR_BASE_VADDR;
        dir->entry[index].present = 1;
        dir->entry[index].rw = 1;
        dir->entry[index].user = 0;
    }

    // 检查内核的大小是否可以放进我们这几个表
    u32 kernel_mem_size = (u32)&__kernel_size_from_start >> PG_SIZE_BITS;
    if (kernel_mem_size > (KERNEL_MAX_TABLE * PG_MAX_ENTRIES)) {
        u16* vga = (u16*)(0xB8000);
        char msg[] = "Kernel too large! > 00MiB";
        u32 size = KERNEL_MAX_TABLE * 4;
        msg[20] += (size / 10) & 0xF;
        msg[21] += (size % 10) & 0xF;
        for (u32 i = 0; i < 26; i++) {
            vga[i] = ((15 << 4 | 12) << 8) | msg[i];
        }
        for (;;);
    }

    // 重映射内核至高半区地址（>=0xC0000000）
    for (u32 i = 0; i < kernel_page_count; i++) {
        u32 tbl_index = i + kernel_tbl_index;

        tbl = (pg_table_t*)(dir->entry[kernel_dir_index + tbl_index / PG_MAX_ENTRIES].data & DIR_BASE_VADDR);
        tbl_index %= PG_MAX_ENTRIES;

        tbl->entry[tbl_index].frame = (kernel_pm >> PG_SIZE_BITS) + i;
        tbl->entry[tbl_index].present = 1;
        tbl->entry[tbl_index].rw = 1;
        tbl->entry[tbl_index].user = 1;
    }

    // 最后一个entry用于循环映射
    dir->entry[PG_MAX_ENTRIES - 1].frame = (ptd & DIR_BASE_VADDR) >> PG_SIZE_BITS;
    dir->entry[PG_MAX_ENTRIES - 1].present = 1;
    dir->entry[PG_MAX_ENTRIES - 1].rw = 1;
    dir->entry[PG_MAX_ENTRIES - 1].pwt = 1;
    dir->entry[PG_MAX_ENTRIES - 1].pcd = 1; // 因为循环映射就是为了修改(频繁)，所以关闭TLB缓存
}

void hhk_init(u32 ptd, u32 k_size) {
    u8* kpg = (u8*)ptd;
    for (u32 i = 0; i < k_size; i++) {
        kpg[i] = 0;
    }
    ptd_init(ptd);
}
